# -*- coding: utf - 8 -*-
"""
KWDB兼容SQL文件生成器

负责将转换后的SQL语句生成为KWDB兼容的SQL文件，并提供批量执行功能。
"""

import os
import subprocess
import logging
from typing import List, Dict, Any, Optional, Tuple
# from pathlib import Path  # noqa: F401

try:
    from .parser import ParsedStatement, StatementType
except ImportError:
    from parser import ParsedStatement, StatementType


class KWDBSQLGenerator:
    """KWDB SQL文件生成器"""

    def __init__(self, config: Optional[Dict[str, Any]] = None):
        self.logger = logging.getLogger(__name__)
        self.config = config or {}

        # 语句执行顺序
        self.execution_order = [
            StatementType.DROP_DATABASE,
            StatementType.CREATE_DATABASE,
            StatementType.USE_DATABASE,
            StatementType.DROP_TABLE,
            StatementType.CREATE_TABLE,
            StatementType.INSERT_INTO,
            StatementType.CREATE_INDEX,
            StatementType.CREATE_SEQUENCE,
        ]

    def generate_sql_file(self, statements: List[ParsedStatement],
                          output_path: str, encoding: str = 'utf - 8') -> bool:
        """
        生成KWDB兼容的SQL文件

        Args:
            statements: 转换后的语句列表
            output_path: 输出文件路径
            encoding: 文件编码

        Returns:
            是否成功生成
        """
        try:
            # 确保输出目录存在
            output_dir = os.path.dirname(output_path)
            if output_dir:
                os.makedirs(output_dir, exist_ok=True)

            # 按执行顺序排序语句
            sorted_statements = self._sort_statements(statements)

            # 生成SQL内容
            sql_content = self._generate_sql_content(sorted_statements)

            # 写入文件
            with open(output_path, 'w', encoding=encoding) as f:
                f.write(sql_content)

            self.logger.info(f"成功生成SQL文件: {output_path}")
            self.logger.info(f"共包含{len(sorted_statements)}个语句")

            return True

        except Exception as e:
            self.logger.error(f"生成SQL文件失败: {e}")
            return False

    def execute_sql_file(self, sql_file_path: str,
                         kwbase_path: str = 'kwbase',
                         host: str = 'localhost',
                         port: int = 26257,
                         database: str = '',
                         insecure: bool = True) -> Tuple[bool, str, str]:
        """
        通过kwbase CLI执行SQL文件

        Args:
            sql_file_path: SQL文件路径
            kwbase_path: kwbase CLI工具路径
            host: KWDB服务器地址
            port: KWDB端口
            database: 数据库名
            insecure: 是否使用非安全连接

        Returns:
            (是否成功, 标准输出, 错误输出)
        """
        try:
            # 构建命令
            cmd = [kwbase_path, 'sql']

            if insecure:
                cmd.append('--insecure')

            cmd.append(f'--host={host}:{port}')

            if database:
                cmd.extend(['--database', database])

            self.logger.info(f"执行命令: {' '.join(cmd)} < {sql_file_path}")

            # 打开SQL文件并通过stdin传递给kwbase
            with open(sql_file_path, 'r', encoding='utf-8') as sql_file:
                # 执行命令，使用stdin输入重定向
                result = subprocess.run(
                    cmd,
                    stdin=sql_file,
                    capture_output=True,
                    text=True,
                    timeout=300  # 5分钟超时
                )

            success = result.returncode == 0

            if success:
                self.logger.info("SQL文件执行成功")
            else:
                self.logger.error(f"SQL文件执行失败，返回码: {result.returncode}")

            return success, result.stdout, result.stderr

        except subprocess.TimeoutExpired:
            self.logger.error("SQL文件执行超时")
            return False, "", "执行超时"
        except Exception as e:
            self.logger.error(f"执行SQL文件时出错: {e}")
            return False, "", str(e)

    def batch_execute_statements(self, statements: List[ParsedStatement],
                                 kwbase_path: str = 'kwbase',
                                 host: str = 'localhost',
                                 port: int = 26257,
                                 database: str = '',
                                 insecure: bool = True,
                                 batch_size: int = 100) -> Dict[str, Any]:
        """
        批量执行语句

        Args:
            statements: 语句列表
            kwbase_path: kwbase CLI工具路径
            host: KWDB服务器地址
            port: KWDB端口
            database: 数据库名
            insecure: 是否使用非安全连接
            batch_size: 批量大小

        Returns:
            执行结果统计
        """
        results = {
            'total': len(statements),
            'success': 0,
            'failed': 0,
            'errors': []
        }

        # 按执行顺序排序
        sorted_statements = self._sort_statements(statements)

        # 分批执行
        for i in range(0, len(sorted_statements), batch_size):
            batch = sorted_statements[i:i + batch_size]

            # 创建临时SQL文件
            temp_file = f"temp_batch_{i//batch_size}.sql"

            try:
                # 生成批次SQL内容
                batch_content = self._generate_sql_content(batch)

                with open(temp_file, 'w', encoding='utf - 8') as f:
                    f.write(batch_content)

                # 执行批次
                success, stdout, stderr = self.execute_sql_file(
                    temp_file, kwbase_path, host, port, database, insecure
                )

                if success:
                    results['success'] += len(batch)
                    self.logger.info(f"批次 {i//batch_size + 1} 执行成功")
                else:
                    results['failed'] += len(batch)
                    results['errors'].append({
                        'batch': i // batch_size + 1,
                        'error': stderr
                    })
                    self.logger.error(f"批次 {i//batch_size + 1} 执行失败: {stderr}")

            except Exception as e:
                results['failed'] += len(batch)
                results['errors'].append({
                    'batch': i // batch_size + 1,
                    'error': str(e)
                })
                self.logger.error(f"批次 {i//batch_size + 1} 处理失败: {e}")

            finally:
                # 清理临时文件
                if os.path.exists(temp_file):
                    os.remove(temp_file)

        return results

    def _sort_statements(self, statements: List[ParsedStatement]) -> List[ParsedStatement]:
        """按执行顺序排序语句"""
        sorted_statements = []

        # 按预定义顺序添加语句
        for stmt_type in self.execution_order:
            type_statements = [s for s in statements if s.statement_type == stmt_type]

            # 对于某些类型的语句，需要进一步排序
            if stmt_type == StatementType.CREATE_TABLE:
                # 表创建顺序可能涉及外键依赖，这里简单按表名排序
                type_statements.sort(key=lambda x: x.table_name or '')
            elif stmt_type == StatementType.INSERT_INTO:
                # INSERT语句按表名分组
                type_statements.sort(key=lambda x: x.table_name or '')

            sorted_statements.extend(type_statements)

        # 添加其他未分类的语句
        other_statements = [s for s in statements if s.statement_type not in self.execution_order]
        sorted_statements.extend(other_statements)

        return sorted_statements

    def _generate_sql_content(self, statements: List[ParsedStatement]) -> str:
        """生成SQL文件内容"""
        lines = []

        # 添加文件头注释
        lines.append("-- MySQL到KWDB迁移生成的SQL文件")
        lines.append(f"-- 生成时间: {self._get_current_timestamp()}")
        lines.append(f"-- 语句总数: {len(statements)}")
        lines.append("")

        current_type = None

        for stmt in statements:
            # 添加语句类型分组注释
            if stmt.statement_type != current_type:
                current_type = stmt.statement_type
                lines.append(f"-- {current_type.value} 语句")
                lines.append("")

            # 添加单个语句注释（可选）
            if stmt.table_name:
                lines.append(f"-- 表: {stmt.table_name}")
            elif stmt.database_name:
                lines.append(f"-- 数据库: {stmt.database_name}")

            # 添加SQL语句
            sql = stmt.original_sql.strip()
            if not sql.endswith(';'):
                sql += ';'

            lines.append(sql)
            lines.append("")

        return '\n'.join(lines)

    def generate_rollback_script(self, statements: List[ParsedStatement],
                                 output_path: str) -> bool:
        """
        生成回滚脚本

        Args:
            statements: 原始语句列表
            output_path: 回滚脚本输出路径

        Returns:
            是否成功生成
        """
        try:
            rollback_statements = []

            # 按相反顺序生成回滚语句
            for stmt in reversed(statements):
                if stmt.statement_type == StatementType.CREATE_TABLE:
                    if stmt.table_name:
                        rollback_sql = f"DROP TABLE IF EXISTS {stmt.table_name};"
                        rollback_statements.append(rollback_sql)

                elif stmt.statement_type == StatementType.CREATE_DATABASE:
                    if stmt.database_name:
                        rollback_sql = f"DROP DATABASE IF EXISTS {stmt.database_name};"
                        rollback_statements.append(rollback_sql)

                elif stmt.statement_type == StatementType.CREATE_INDEX:
                    # 从原始SQL中提取索引名
                    index_name = self._extract_index_name(stmt.original_sql)
                    if index_name:
                        rollback_sql = f"DROP INDEX IF EXISTS {index_name};"
                        rollback_statements.append(rollback_sql)

            # 生成回滚脚本内容
            lines = [
                "-- 回滚脚本",
                f"-- 生成时间: {self._get_current_timestamp()}",
                "",
                "-- 警告: 此脚本将删除迁移过程中创建的所有对象",
                "-- 请谨慎使用！",
                ""
            ]

            lines.extend(rollback_statements)

            with open(output_path, 'w', encoding='utf - 8') as f:
                f.write('\n'.join(lines))

            self.logger.info(f"成功生成回滚脚本: {output_path}")
            return True

        except Exception as e:
            self.logger.error(f"生成回滚脚本失败: {e}")
            return False

    def _extract_index_name(self, sql: str) -> Optional[str]:
        """从CREATE INDEX语句中提取索引名"""
        import re

        # 匹配 CREATE [UNIQUE] INDEX index_name
        match = re.search(r'CREATE\s+(?:UNIQUE\s+)?INDEX\s+(\w+)', sql, re.IGNORECASE)
        if match:
            return match.group(1)

        return None

    def _get_current_timestamp(self) -> str:
        """获取当前时间戳"""
        from datetime import datetime
        return datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    def generate_migration_report(self, original_statements: List[ParsedStatement],
                                  converted_statements: List[ParsedStatement],
                                  output_path: str) -> bool:
        """
        生成迁移报告

        Args:
            original_statements: 原始语句列表
            converted_statements: 转换后语句列表
            output_path: 报告输出路径

        Returns:
            是否成功生成
        """
        try:
            report_lines = []

            # 报告头
            report_lines.extend([
                "# MySQL到KWDB迁移报告",
                f"生成时间: {self._get_current_timestamp()}",
                "",
                "## 迁移统计",
                f"- 原始语句数量: {len(original_statements)}",
                f"- 转换后语句数量: {len(converted_statements)}",
                f"- 转换成功率: {len(converted_statements)/len(original_statements)*100:.1f}%",
                ""
            ])

            # 按类型统计
            report_lines.append("## 按语句类型统计")
            report_lines.append("")

            type_stats = {}
            for stmt in original_statements:
                stmt_type = stmt.statement_type.value
                type_stats[stmt_type] = type_stats.get(stmt_type, 0) + 1

            converted_type_stats = {}
            for stmt in converted_statements:
                stmt_type = stmt.statement_type.value
                converted_type_stats[stmt_type] = converted_type_stats.get(stmt_type, 0) + 1

            for stmt_type, original_count in type_stats.items():
                converted_count = converted_type_stats.get(stmt_type, 0)
                success_rate = converted_count / original_count * 100 if original_count > 0 else 0

                report_lines.append(f"- {stmt_type}: {original_count} -> {converted_count} ({success_rate:.1f}%)")

            report_lines.append("")

            # 表信息统计
            tables = set()
            for stmt in original_statements:
                if stmt.table_name:
                    tables.add(stmt.table_name)

            if tables:
                report_lines.extend([
                    "## 涉及的表",
                    f"共 {len(tables)} 个表:",
                    ""
                ])

                for table in sorted(tables):
                    report_lines.append(f"- {table}")

                report_lines.append("")

            # 数据库信息
            databases = set()
            for stmt in original_statements:
                if stmt.database_name:
                    databases.add(stmt.database_name)

            if databases:
                report_lines.extend([
                    "## 涉及的数据库",
                    ""
                ])

                for db in sorted(databases):
                    report_lines.append(f"- {db}")

                report_lines.append("")

            # 注意事项
            report_lines.extend([
                "## 注意事项",
                "",
                "1. 请在执行前备份目标数据库",
                "2. 建议先在测试环境验证迁移结果",
                "3. 检查生成的SQL文件是否符合预期",
                "4. 关注错误日志中的警告信息",
                ""
            ])

            # 写入报告文件
            with open(output_path, 'w', encoding='utf - 8') as f:
                f.write('\n'.join(report_lines))

            self.logger.info(f"成功生成迁移报告: {output_path}")
            return True

        except Exception as e:
            self.logger.error(f"生成迁移报告失败: {e}")
            return False
